import os
import time
import logging
from selenium import webdriver
from selenium.webdriver.edge.service import Service
from selenium.webdriver.edge.options import Options
from urllib.parse import urljoin
import requests
from bs4 import BeautifulSoup


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def scroll_down(driver, num_scrolls=10, scroll_pause_time=2, max_attempts=5):

    last_height = driver.execute_script("return document.body.scrollHeight")
    attempts = 0
    
    for _ in range(num_scrolls):

        driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")

        time.sleep(scroll_pause_time)

        new_height = driver.execute_script("return document.body.scrollHeight")
        if new_height == last_height:
            attempts += 1
            if attempts >= max_attempts:
                break
        else:
            attempts = 0
        last_height = new_height

def download_images(url, save_folder):

    if not save_folder:
        logging.error("Save path cannot be empty, please re-enter")
        return

    if not os.path.exists(save_folder):
        os.makedirs(save_folder)


    edge_options = Options()
    edge_options.add_argument("--headless")  
    edge_options.add_argument("--disable-gpu")
    edge_options.add_argument("--no-sandbox")


    service = Service("D:\Program Files\MEWebdriver\msedgedriver.exe")  # 替换为你的msedgedriver路径


    driver = None
    try:
        driver = webdriver.Edge(service=service, options=edge_options)
        driver.get(url)

        scroll_down(driver, num_scrolls=60, scroll_pause_time=3)  # 增加滚动次数

        page_source = driver.page_source

        soup = BeautifulSoup(page_source, 'html.parser')

        img_tags = soup.find_all('img')
        img_urls = [urljoin(url, img['src']) for img in img_tags if 'src' in img.attrs]

        logging.info(f"Parsed image link: {img_urls}")

        for idx, img_url in enumerate(img_urls):
            try:
                response = requests.get(img_url, proxies=None, timeout=10)
                if response.status_code == 200:

                    img_name = f"image_{idx}.png"
                    img_path = os.path.join(save_folder, img_name)
                    with open(img_path, 'wb') as f:
                        f.write(response.content)
                    logging.info(f"sucessful: {img_url} -> {img_path}")
                else:
                    logging.warning(f"Unable to download image: {img_url},status: {response.status_code}")
            except requests.exceptions.RequestException as e:
                logging.error(f"An error occurred while downloading the image: {img_url} - {e}")
            except Exception as e:
                logging.error(f"unknown error: {img_url} - {e}")
            

            time.sleep(1) 
    except Exception as e:
        logging.error(f"An error occurred: {e}")
    finally:
        if driver:
            driver.quit()

if __name__ == "__main__":

    url = input("Please enter the website address: ")
    save_folder = input("Please enter the folder path to save the image: ")

    download_images(url, save_folder)
